from pvz_env import PVZEnv
from dqn_agent import DQNAgent, train
import os

def main():
    # 创建游戏环境
    env = PVZEnv()
    
    # 创建DQN智能体
    agent = DQNAgent(
        env,
        learning_rate=1e-4,
        gamma=0.99,
        epsilon_start=1.0,
        epsilon_final=0.01,
        epsilon_decay=0.995,
        memory_size=10000,
        batch_size=32,
        target_update=1000
    )
    
    # 检查是否存在已保存的模型
    model_path = 'pvz_model.pth'
    flag=input("是否使用已保存的模型？(y/n)")
    if os.path.exists(model_path) and flag=='y':
        print("加载已保存的模型...")
        agent.load_model(model_path)
        # 应用模式：使用训练好的模型进行游戏
        play_game(env, agent)
    else:
        # 训练模式
        print("开始训练 DQN 智能体...")
        episode_rewards = train(env, agent, num_episodes=1000)
        print("训练完成！")
        
        # 保存训练好的模型
        print("保存模型...")
        agent.save_model(model_path)

def play_game(env, agent):
    """使用训练好的模型进行游戏"""
    print("开始游戏...")
    state = env.reset()
    done = False
    total_reward = 0
    
    while not done:
        # 使用训练好的模型选择动作
        action = agent.select_action(state)
        next_state, reward, done, _ = env.step(action)
        
        total_reward += reward
        state = next_state
        
        # 渲染游戏状态
        env.render()
    
    print(f"游戏结束！总得分：{total_reward}")

if __name__ == '__main__':
    main()
